#include "InjectKit.h"

#include <stdio.h>
#include <WinBase.h>
#include <TlHelp32.h>
#include "Psapi.h"
#include <iphlpapi.h>

#pragma comment(lib,"Psapi.lib")

#include "../hookCommon.h"


InjectKit::InjectKit(char* target)
{
	clean();

	targetName = tolowerstr(target);
	init(targetName);
}

InjectKit::~InjectKit()
{
	if(targetHandle)
		CloseHandle(targetHandle);

	if(mapMemory)
		UnmapViewOfFile(mapMemory);

	if(mapHandle)
		CloseHandle(mapHandle);

	if(targetName)
		free(targetName);

	for(size_t i = 0 ; i < inject_dlls.size(); i ++)
		free(inject_dlls[i]);

	clean();
}

void InjectKit::clean()
{
	targetHandle = 0;
	targetProcessID = 0;
	mapHandle = 0;
	mapMemory = 0;
	targetName = 0;

	inject_dlls.clear();
}

int InjectKit::init(char* target)
{
	HANDLE SnapShot_h = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS,NULL);
	if( NULL  == SnapShot_h )
		return InjectError("CreateToolhelp32Snapshot failed!");

	PROCESSENTRY32  Entry32; 
	Entry32.dwSize = sizeof(Entry32);

	if( FALSE == Process32First(SnapShot_h,&Entry32) )
		return InjectError("Process32First failed!");

	while(Process32Next(SnapShot_h,&Entry32))
	{
		char* str1 = tolowerstr(Entry32.szExeFile);
		char* str2 = tolowerstr(target);

		if( 0 == strcmp(str1,str2))
		{
			free(str1);
			free(str2);

			HANDLE handle = OpenProcess(PROCESS_ALL_ACCESS,FALSE,Entry32.th32ProcessID);
			if( NULL == handle )
				return InjectError("OpenProcess failed!");

			targetProcessID = Entry32.th32ProcessID;
			targetHandle = handle;
			CloseHandle(SnapShot_h);
			return 0;
		}

		free(str1);
		free(str2);
	}

	CloseHandle(SnapShot_h);
	return InjectError("Not Found!");
}


HMODULE InjectKit::getModule(char* module)
{
	if(module == 0)
		module = targetName;

	HMODULE  hMod[1024]={0};    
	DWORD len = 0;
	EnumProcessModules(targetHandle,hMod,sizeof(hMod),&len);
	DWORD times = len/sizeof(DWORD);

	for(DWORD i = 0; i < times ; i++)
	{
		char name[1024]={0};
		GetModuleBaseName(targetHandle,hMod[i],name,sizeof(name));
		char* str1 =tolowerstr(name);
		char* str2 =tolowerstr(module);
		int result = strcmp(str1,str2);
		free(str1);
		free(str2);
		if(result == 0)
			return hMod[i];
	}
	return 0;
}

int InjectKit::load_dll(char* dllname)
{
	char dllPath[MAX_PATH]={0};

	if(0 ==	GetFullPathName(dllname,sizeof(dllPath),dllPath,NULL))
		return InjectError("LoadDll GetFullPathName failed!");

	LPVOID  MemorySpace = VirtualAllocEx(targetHandle,NULL,sizeof(dllPath),MEM_COMMIT,PAGE_EXECUTE_READWRITE);
	if( NULL == MemorySpace)
		return InjectError("LoadDll CreateVirtualMemory failed!");

	if( 0 == WriteProcessMemory(targetHandle,MemorySpace,dllPath,sizeof(dllPath),NULL))
		return InjectError("LoadDll WriteMemoryError");

	HMODULE kernel32 = GetModuleHandle("Kernel32");
	if( NULL == kernel32)
		return InjectError("LoadDll GetKernel32 ModuleHandle failed");

	HANDLE crt_h = CreateRemoteThread(targetHandle,0,0,(LPTHREAD_START_ROUTINE)GetProcAddress(kernel32,"LoadLibraryA"),MemorySpace,0,NULL);
	if( NULL == crt_h)
		return InjectError("LoadDll createRemoteThread failed");

	WaitForSingleObject(crt_h,INFINITE);
	VirtualFreeEx(targetHandle,MemorySpace,sizeof(dllPath),MEM_RELEASE);
	CloseHandle(kernel32);
	CloseHandle(crt_h);

	size_t size = strlen(dllname);
	char* hookDllName = (char*)malloc(sizeof(char)*size + 1);
	memcpy(hookDllName,dllname,sizeof(char)*size);
	hookDllName[size]=0;

	inject_dlls.push_back(hookDllName);

	return 0;
}

int InjectKit::unload_dll(char* dllname)
{
	if(dllname == 0 && inject_dlls.size() != 1)
		return  InjectError("Not set unload dll name!");

	if(dllname == 0)
		dllname = inject_dlls[0];

	size_t size  = strlen(dllname);
	char* DLLNAME = (char*)malloc(sizeof(char*)*size + 1);
	memcpy(DLLNAME,dllname,sizeof(char)*size);
	DLLNAME[size] = 0;

	LPVOID dllAlloc = VirtualAllocEx(targetHandle,NULL,size + 1 ,MEM_COMMIT,PAGE_EXECUTE_READWRITE);
	if( NULL == dllAlloc)
		return InjectError("UnLoadDll VirtualAlloc failed!");

	BOOL result = FALSE;
	result = WriteProcessMemory(targetHandle,dllAlloc,DLLNAME, size + 1 ,NULL);
	if(result == FALSE)
		return InjectError("UnLoadDll Write failed!");

	free(DLLNAME);

	HMODULE kernel32 = GetModuleHandle("Kernel32");

	HANDLE th = CreateRemoteThread(targetHandle,NULL,0,
		(LPTHREAD_START_ROUTINE)GetProcAddress(kernel32,"GetModuleHandleA"),dllAlloc,0,0);

	if(th == NULL)
		return InjectError("Unload GetModuleHandle failed!");

	WaitForSingleObject(th,INFINITE);

	HMODULE dll_M = 0;
	GetExitCodeThread(th,(LPDWORD)&dll_M);
	CloseHandle(th);
	VirtualFreeEx(targetHandle,dllAlloc,0,MEM_RELEASE);

	th  = CreateRemoteThread(targetHandle,NULL,0,
		(LPTHREAD_START_ROUTINE)GetProcAddress(kernel32,"FreeLibrary"),dll_M,0,0);

	WaitForSingleObject(th,INFINITE);

	if(th == NULL)
		return InjectError("Unload freelibrary failed!");

	CloseHandle(kernel32);
	CloseHandle(th);
	return 0;
}

int	InjectKit::CreateMapMemory(int size)
{
	mapHandle = CreateFileMapping(INVALID_HANDLE_VALUE, NULL , 
		PAGE_EXECUTE_READWRITE | SEC_COMMIT , 0 ,size,(LPCSTR)INJECT_MEMORY_ID);	

	if(NULL == mapHandle)
		return InjectError("CreateFileMapping  failed!");

	mapMemory =MapViewOfFile(mapHandle,FILE_MAP_ALL_ACCESS,0,0,0);

	if(NULL == mapMemory)
		return InjectError("MapVieOfFile Error!");

	memset(mapMemory,0,size);
	return 0;	
}


int InjectKit::InjectError(char* message)
{
	int code = ::GetLastError();
	char error_tips[1024] = {0};
	sprintf(error_tips,"%s Error:%d",message,code);
	MessageBoxA(NULL,message,"ERROR",MB_OK);
	return code;
}

